package my.edu.ai.controller;

import cn.hutool.json.JSONObject;

import jakarta.validation.constraints.NotNull;
import my.edu.ai.service.ChatMemoryServiceImpl;
import my.edu.ai.service.IAiChatService;
import my.edu.common.domain.Response;
import my.edu.common.utils.RedisUtils;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.core.scheduler.Schedulers;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;

@RestController
@RequestMapping("/ai")
@Validated
public class AIController {

    //对话客户端
    @Autowired
    private ChatClient chatClient;
    @Autowired
    private ChatModel chatModel;
    @Autowired
    private IAiChatService aiChatService;
    @Autowired
    private ChatMemory chatMemory = new ChatMemoryServiceImpl();
    @Autowired
    private RedisUtils redisUtils;

    //聊天停止器
    private final ConcurrentHashMap<String, Sinks.One<Boolean>> stopSignalMap = new ConcurrentHashMap<>();
    //历史记录
    private final ConcurrentHashMap<String, List<String>> historyMap = new ConcurrentHashMap<>();


    @PostMapping("/newChat")
    public Response newChat() {
        return Response.success(aiChatService.addChat());
    }


    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> stream(
            @NotNull String message,
            @NotNull Integer chatId) {
        String conversationId = aiChatService.getConversationId(chatId);
        MessageChatMemoryAdvisor advisor = new MessageChatMemoryAdvisor(chatMemory, conversationId, 10);
        stopSignalMap.put(conversationId, Sinks.one());
        // 初始化历史记录存储
        historyMap.put(conversationId, new ArrayList<>());

        Flux<String> aiStream = chatClient.prompt()
                .user(message)
                .advisors(advisor)
                .stream()
                .content()
                .map(text -> {
                    historyMap.get(conversationId).add(text); // 记录 AI 响应内容
                    return wrapResponse(chatId, text);
                })
                // 添加响应式生命周期管理
                .doOnSubscribe(subscription -> {
                    System.out.println("客户端开始订阅流，chatId: " + chatId);
//                    redisUtils.set(conversationId + ":stream", "start");
                })
                .takeUntilOther(stopSignalMap.get(conversationId).asMono())  // 监听停止信号
                .doOnCancel(() -> {
                    System.out.println("客户端主动断开，chatId: " + chatId);
//                    redisUtils.set(conversationId + ":stream", "end");
                    // 目前未出现
                })
                .doFinally(signal -> {
                    System.out.println("流处理结束，原因：" + signal + "，chatId: " + chatId);
//                    redisUtils.set(conversationId + ":stream", "end");
                });

        Flux<String> endSignalFlux = Flux.just("123END123")
                .map(text ->{
                    return wrapResponse(chatId, text);
                });

        return aiStream.concatWith(endSignalFlux);
    }

    // 停止某个 AI 对话
    @RequestMapping(value = "/stop", method = {RequestMethod.GET, RequestMethod.POST})
    public Mono<List<String>> stopChat(@RequestParam Integer chatId) {
        String conversationId = aiChatService.getConversationId(chatId);
        Sinks.One<Boolean> stopSignal = stopSignalMap.get(conversationId);

        if (stopSignal != null) {
            stopSignal.tryEmitValue(true); // 触发停止信号
            stopSignalMap.remove(conversationId); // 移除停止信号
//            redisUtils.set(conversationId + ":stream", "end");
            // 获取已经生成的 AI 内容
            List<String> history = historyMap.getOrDefault(conversationId, new ArrayList<>());
            historyMap.remove(conversationId); // 清理历史数据
            chatMemory.add(conversationId, new AssistantMessage(String.join("", history)));
            return Mono.just(history);
        }

        return Mono.just(new ArrayList<>()); // 如果没有找到聊天，返回空列表
    }
    private String wrapResponse(Integer chatId, String text) {
        JSONObject json = new JSONObject();
        json.put("chatId", chatId);
        json.put("textContent", text);
        return json.toString();
    }

    @DeleteMapping
    public Response delChat(@NotNull(message = "chatId不能为空") Integer chatId) {
        aiChatService.removeByChatId(chatId);
        return Response.success();
    }

    @GetMapping("/list")
    public Response list() {
        return Response.success(aiChatService.listChat());
    }

    @GetMapping("/one")
    public Response getOne(@NotNull(message = "chatId不能为空") Integer chatId) {
        return Response.success(aiChatService.getOneChat(chatId));
    }

//    @GetMapping("/msg")
//    Response history(@NotNull(message = "消息不能为空") String message, @NotNull(message = "chatId不能为空") Integer chatId) {
//        String conversationId = aiChatService.getConversationId(chatId);
//        MessageChatMemoryAdvisor messageChatMemoryAdvisor = new MessageChatMemoryAdvisor(chatMemory, conversationId, 10);
//        String content = chatClient.prompt()
//                //用户信息
//                .user(message)
//                .advisors(messageChatMemoryAdvisor)
//                //远程请求大模型
//                .call()
//                //返回文本
//                .content();
//        return Response.success("", new JSONObject().putOnce("textContent", content).putOnce("chatId", chatId));
//    }



//    @GetMapping(value = "/s", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
//    public Flux<String> stream(@NotNull(message = "消息不能为空") String message) {
//        ChatMemory chatMemory2 = new InMemoryChatMemory();
//        MessageChatMemoryAdvisor messageChatMemoryAdvisor = new MessageChatMemoryAdvisor(chatMemory2, 1+"", 10);
//        // 使用 Flux 流式处理，并包装 chatId 和内容
//        Flux<String> stringFlux = chatClient.prompt()
//                .user(message)
//                .advisors(messageChatMemoryAdvisor)
//                .stream()
//                .content()
//                .publishOn(Schedulers.boundedElastic());
//
//        return    stringFlux.concatWith(Flux.just("123END123"));
//    }
//
//
//    @GetMapping("/ai")
//    String generation(@RequestParam(required = false, defaultValue = "你好！") String message) {
//        //prompt：提示词
//        return this.chatClient.prompt()
//                //用户信息
//                .user(message)
//                //远程请求大模型
//                .call()
//                //返回文本
//                .content();
//    }
//
//
//    @GetMapping("/chatModel")
//    String chatModel(String message) {
//        ChatResponse response = chatModel.call(
//                new Prompt(
//                        message,
//                        OpenAiChatOptions.builder()
//                                .model("gpt-4-o")
//                                .temperature(0.4)
//                                .build()
//                ));
//        return response.getResult().getOutput().getContent();
//    }
}